
from  munch import Munch
import yaml
from abc import abstractmethod, ABCMeta
from fvcore.common.config import CfgNode as _CfgNode
from argparse import ArgumentParser, RawDescriptionHelpFormatter

__all__ = ['CfgNode', 'get_config', 'ArgsParserCallback', 'C']

print('__name__:', __name__)


class CfgNode(Munch):
    def __str__(self):
        return str(self.toDict())


class ArgsParserCallback:
    @abstractmethod
    def add_parser(self, parser: ArgumentParser):
        # raise NotImplementedError('function add_parser not impl.')
        pass


class ArgsParser(ArgumentParser):
    def __init__(self, callback: ArgsParserCallback=None):
        super(ArgsParser, self).__init__(
            formatter_class=RawDescriptionHelpFormatter)
        self.add_argument("-c", "--config", help="configuration file to use")
        self.add_argument(
            "-o", "--opt", nargs='+', help="set configuration options")
        if callback is not None:
            callback.add_parser(self)

    def parse_args(self, argv=None):
        args = super(ArgsParser, self).parse_args(argv)
        assert args.config is not None, \
            "Please specify --config=configure_file_path."
        args.opt = self._parse_opt(args.opt)
        return args

    def _parse_opt(self, opts):
        config = {}
        if not opts:
            return config
        for s in opts:
            s = s.strip()
            k, v = s.split('=')
            config[k] = yaml.load(v, Loader=yaml.Loader)
        return config


def _update_cfg_from_yaml(config, cfg, root):
    for k, v in cfg.items():
        if isinstance(v, dict):
            config[k] = CfgNode()
            _update_cfg_from_yaml(config, v, config[k])
        else:
            root[k] = v


def merge_config_from_opt(opt, config):
    """
    Merge opt into config.
    Args:
        opt (dict): opts.
        config (Munch): opt to be merged.
    Returns: config
    """
    for key, value in opt.items():
        if "." not in key:
            if isinstance(value, dict) and key in config:
                config[key].update(value)
            else:
                config[key] = value
        else:
            sub_keys = key.split('.')
            assert (
                sub_keys[0] in config
            ), "the sub_keys can only be one of config: {}, but get: {}, please check your running command".format(
                config.keys(), sub_keys[0])
            cur = config[sub_keys[0]]
            for idx, sub_key in enumerate(sub_keys[1:]):
                assert (
                    sub_key in cur
                ), "the sub_keys can only be one of config: {}, but get: {}, please check your running command".format(
                    cur.keys(), sub_key)
                if idx == len(sub_keys) - 2:
                    cur[sub_key] = value
                else:
                    cur = cur[sub_key]


C = CfgNode()

def get_config(callback: ArgsParserCallback=None, def_config: CfgNode=None):
    assert isinstance(def_config, CfgNode)

    FLAGS = ArgsParser(callback=callback).parse_args()
    config = C
    if def_config is not None:
        config.update(def_config)
    cfg = _CfgNode.load_yaml_with_base(FLAGS.config, allow_unsafe=True)
    _update_cfg_from_yaml(config, cfg, config)
    merge_config_from_opt(FLAGS.opt, config)
    return C
